# -*- coding: utf-8 -*-
"""
文件名: prompt_constructor.py
功能:
  - 构造多模态问答对（支持: 仅质量 / 仅疾病 / 质量+疾病 / 选择题）
  - 严格按“单病/多病”分别生成 判断题/比较题/诊断题/选择题
  - 引入疾病关系映射(diease_relation_map)，控制负样本采样:
        * black_list: 禁止作为负样本
        * white_list: 优先作为强负样本
        * 若冲突(某疾病出现在黑白名单)，黑名单优先
  - 支持强负样本与弱负样本混合采样，防止模式单一
  - 选择题（A/B/C/D 选项）。单病→单选；多病(<4)→多选；多病数≥4时不构造选择题
  - 返回字典结构:
        {"judge":[], "compare":[], "diagnosis":[], "quality":[], "quality+diagnosis":[], "choice":[]}
  - 所有异常立即 raise，并输出关键信息(疾病名、采样池状态等)
"""

import random
from typing import List, Tuple, Optional, Sequence, Dict
import json


class PromptConstructor:
    def __init__(self, all_cn_diseases: List[str], seed: Optional[int] = 42, diease_relation_map: Optional[str] = None):
        """
        参数:
            all_cn_diseases: 系统内支持的中文疾病名称全集(>0)
            seed: 随机种子(可复现)
            diease_relation_map: 疾病关系映射文件路径(JSON)，包含黑白名单
        """
        if not isinstance(all_cn_diseases, Sequence) or len(all_cn_diseases) == 0:
            raise ValueError("all_cn_diseases 为空或类型不正确，需为包含至少 1 个疾病名的序列。")

        self.all_cn_diseases = list(all_cn_diseases)
        self.rng = random.Random(seed)

        # 载入疾病关系映射
        self.diease_relation_map = {}
        if diease_relation_map:
            try:
                with open(diease_relation_map, 'r', encoding="utf-8") as f:
                    self.diease_relation_map = json.load(f)
            except Exception as e:
                raise RuntimeError(f"加载疾病关系映射失败: {diease_relation_map}, 错误信息: {e}")

        # ---------------- 模板池定义 ---------------- #
        self.judge_q_tpls = [
            "图像中显示的是{}吗？请回答是或否。",
            "该眼底照片能诊断为{}吗？请回答是或否。",
            "是否可以从该图像诊断出{}？请回答是或否。",
            "患者是否患有{}？请回答是或否。",
            "请判断图像所示是否为{}，请回答是或否。"
        ]
        self.judge_a_true_single = [
            "是。",
            "是，可以观察到{}的特征。",
            "是，符合{}的诊断要点。"
        ]
        self.judge_a_true_multi = [
            "是，除{}外，还可见{}等病变特征。",
            "是，可见{}表现，同时伴随{}。",
            "是，存在{}征象，并伴随{}。"
        ]
        self.judge_a_false = [
            "否，未见{}的典型征象。",
            "否，图像中未观察到{}相关特征。",
            "否，不支持{}的诊断。"
        ]

        self.diagnosis_q_tpls = [
            "请根据图像直接给出诊断结论。",
            "该患者的眼底情况如何？请直接诊断。",
            "根据图像，你认为患者的诊断结果是什么？",
            "请直接指出本图像的疾病诊断结果。",
            "分析眼底照片后，给出你的诊断结论。"
        ]
        self.diagnosis_a_single = [
            "图像中可见{}的典型表现，诊断为{}。",
            "影像特征符合{}，诊断为{}。",
            "整体表现与{}一致，诊断为{}。"
        ]
        self.diagnosis_a_multi = [
            "图像显示多种病变特征，包括{}。",
            "可见{}等多病变表现，诊断为{}并存。",
            "观察到{}等共现征象。"
        ]

        self.compare_q_tpls = [
            "图像更可能属于以下哪种疾病：{}或{}？请直接给出疾病名称。",
            "患者的眼底情况更符合{}还是{}？请直接回答疾病名称。",
            "根据图像，{}和{}中哪种疾病的可能性更高？请直接给出疾病名称。"
        ]
        self.compare_a_prefer = [
            "{}的可能性更高。",
            "图像更接近{}。",
            "图像显示{}特征更为明显。"
        ]
        self.compare_a_both_true = [
            "两种病变均可观察到{}特征。",
            "图像显示二者均存在。",
            "两者均可见。"
        ]
        self.compare_a_both_false = [
            "图像未见两者的典型征象。",
            "当前图像不支持{}或{}的诊断。",
            "两种疾病的特征在图像中均不明显。"
        ]

        self.quality_q_tpls = [
            "请对这张眼底图像的质量进行评估。",
            "请描述这张眼底图像的质量情况。",
            "请给出该图像的质量评价，包括清晰度、取景、伪影等。",
            "请判断该图像质量是否足以进行疾病诊断。",
            "请根据图像质量给出分析。"
        ]

        # ---- 新增：选择题模板（不提示单/多选） ---- #
        self.choice_q_tpls = [
            "从以下选项中选择最符合该图像的疾病：\n{options}\n请直接回答字母。",
            "参考下列选项，选择与图像最匹配的疾病：\n{options}\n请直接回答字母。",
            "请选择与该眼底图像相符的疾病：\n{options}\n请直接回答字母。"
        ]
        # 选择题答案格式：字母优先，可附疾病名（更易对齐/调试）
        # 单选示例： "A {糖尿病视网膜病变}"
        # 多选示例： "A {黄斑裂孔}, C {近视盘旁萎缩}"

    # ========================= 对外主接口 ========================= #

    def construct_prompts(
        self,
        true_diag: Optional[List[str]] = None,
        quality_text: Optional[str] = None,
        normal_key: str = "正常眼底",
        normal_fvf_prob: float = 0.1,
        build_choice: bool = True   # 新增：是否构造选择题
    ) -> Dict[str, List[Tuple[str, str]]]:
        """
        返回:
            dict[str, list[tuple(question, answer)]]
        """
        result: Dict[str, List[Tuple[str, str]]] = {
            "quality": [], "judge": [], "compare": [], "diagnosis": [], "quality+diagnosis": [], "choice": []
        }

        
        if quality_text :
            q = self._rc(self.quality_q_tpls)
            a = f"图像质量评估如下：{quality_text}"
            result["quality"].append((q, a))
            # --- 仅质量 --- #
            if (not true_diag or len(true_diag) == 0):
                return result

        # --- 疾病相关 --- #
        if true_diag and len(true_diag) > 0:
            t_list = self._dedup_preserve(true_diag)
            self._validate_true_in_all(t_list)

            # 单病
            if len(t_list) == 1:
                t = t_list[0]
                # 采 4 个负样本，便于后续打乱与多处使用（修复你原代码 TODO）
                f_list = self._sample_negatives_with_relation([t], need=4)
                if len(f_list) < 3:
                    raise ValueError(f"负样本不足以构造题目（需要≥3），target={t}, got={f_list}")
                self.rng.shuffle(f_list)
                f1, f2, f3 = f_list[0], f_list[1], f_list[2]

                # 判断
                result["judge"].append(self._build_judge_positive_single(t))
                result["judge"].append(self._build_judge_negative(f1))

                # 比较：两道不同的负样本，避免 f2 被复用两次
                result["compare"].append(self._build_compare_prefer(t, f2, prefer=t))
                result["compare"].append(self._build_compare_prefer(f3, t, prefer=t))

                # 两者皆否
                if t == normal_key:
                    if self.rng.random() < max(0.0, min(1.0, normal_fvf_prob)):
                        # 再取一个负样本用于 both_false
                        f4 = f_list[3] if len(f_list) >= 4 else f1
                        result["compare"].append(self._build_compare_both_false(f3, f4))
                else:
                    f4 = f_list[3] if len(f_list) >= 4 else f1
                    result["compare"].append(self._build_compare_both_false(f3, f4))

                # 诊断
                result["diagnosis"].append(self._build_diagnosis_single(t))

                # 选择题（单选）
                if build_choice:
                    if len(f_list) < 3:
                        raise ValueError(f"选择题负样本不足（need=3），target={t}, got={f_list}")
                    choice_q, choice_a = self._build_choice_single(t, f_list[:3])
                    result["choice"].append((choice_q, choice_a))

            # 多病
            else:
                # 负样本数与正类数对齐（用于比较题）
                f_list = self._sample_negatives_with_relation(t_list, need=len(t_list))
                if len(f_list) < len(t_list):
                    raise ValueError(f"负样本不足以构造比较题，targets={t_list}, got={f_list}")

                # 判断（多病）
                for t in t_list:
                    result["judge"].append(self._build_judge_positive_multi(t, [x for x in t_list if x != t]))
                for f in f_list:
                    result["judge"].append(self._build_judge_negative(f))

                # 比较：两者皆真
                if len(t_list) >= 2:
                    t1, t2 = self._sample_two_from_list(t_list)
                    result["compare"].append(self._build_compare_both_true(t1, t2, prefer=self._rc([t1, t2])))

                # 比较：t vs f（配对）
                for t, f in zip(t_list, f_list):
                    result["compare"].append(self._build_compare_prefer(t, f, prefer=t))

                # 比较：两者皆否（从全集再采 2 个不在正类里的负样本）
                remaining_space = len(self.all_cn_diseases) - len(set(t_list))
                if remaining_space >= 2:
                    ff = self._sample_negatives_with_relation(t_list, need=2)
                    if len(ff) == 2:
                        result["compare"].append(self._build_compare_both_false(ff[0], ff[1]))

                # 诊断（多病）
                result["diagnosis"].append(self._build_diagnosis_multi(t_list))

                # 选择题（多选，正类数<4 才构造；否则跳过）
                if build_choice and len(t_list) < 4:
                    need_neg = 4 - len(t_list)
                    if need_neg > 0:
                        cand_negs = self._sample_negatives_with_relation(t_list, need=need_neg)
                        if len(cand_negs) < need_neg:
                            raise ValueError(f"选择题负样本不足（need={need_neg}），targets={t_list}, got={cand_negs}")
                    else:
                        cand_negs = []
                    choice_q, choice_a = self._build_choice_multi(t_list, cand_negs)
                    result["choice"].append((choice_q, choice_a))

            # --- 联合质量 --- #
            if quality_text:
                q_quality = self._rc([
                    "请先评估该图像的质量情况。",
                    "首先，请对图像质量进行分析。",
                    "请说明该眼底照片的质量状况。"
                ])
                a_quality = f"图像质量评估如下：{quality_text}"
                q_diag, a_diag = random.choice(result["diagnosis"])
                joint_q = f"{q_quality}\n接着，请你进行诊断：{q_diag}"
                joint_a = f"{a_quality}\n根据图像内容，{a_diag}"
                result["quality+diagnosis"].append((joint_q, joint_a))

        return result # { "quality": [...], "judge": [...], "compare": [...], "diagnosis": [...], "quality+diagnosis": [...], "choice": [..] }

    # ========================= 内部构造子程序 ========================= #

    # --- 判断题 ---
    def _build_judge_positive_single(self, t: str) -> Tuple[str, str]:
        q = self._rc(self.judge_q_tpls).format(t)
        a = self._rc(self.judge_a_true_single).format(t)
        return q, a

    def _build_judge_positive_multi(self, t: str, others: List[str]) -> Tuple[str, str]:
        q = self._rc(self.judge_q_tpls).format(t)
        if not others:
            a = self._rc(self.judge_a_true_single).format(t)
        else:
            others_str = "、".join(others)
            a = self._rc(self.judge_a_true_multi).format(t, others_str)
        return q, a

    def _build_judge_negative(self, f: str) -> Tuple[str, str]:
        q = self._rc(self.judge_q_tpls).format(f)
        a = self._rc(self.judge_a_false).format(f)
        return q, a

    # --- 比较题 ---
    def _build_compare_prefer(self, a: str, b: str, prefer: str) -> Tuple[str, str]:
        q = self._rc(self.compare_q_tpls).format(a, b)
        a_text = self._rc(self.compare_a_prefer).format(prefer)
        return q, a_text

    def _build_compare_both_true(self, a: str, b: str, prefer: str) -> Tuple[str, str]:
        q = self._rc(self.compare_q_tpls).format(a, b)
        a_text = self._rc(self.compare_a_both_true).format(prefer)
        return q, a_text

    def _build_compare_both_false(self, a: str, b: str) -> Tuple[str, str]:
        q = self._rc(self.compare_q_tpls).format(a, b)
        a_text = self._rc(self.compare_a_both_false).format(a, b)
        return q, a_text

    # --- 诊断题 ---
    def _build_diagnosis_single(self, t: str) -> Tuple[str, str]:
        q = self._rc(self.diagnosis_q_tpls)
        a = self._rc(self.diagnosis_a_single).format(t, t)
        return q, a

    def _build_diagnosis_multi(self, t_list: List[str]) -> Tuple[str, str]:
        q = self._rc(self.diagnosis_q_tpls)
        joined = "、".join(t_list)
        tpl = self._rc(self.diagnosis_a_multi)
        count = tpl.count("{}")
        if count == 1:
            a = tpl.format(joined)
        elif count >= 2:
            a = tpl.format(joined, joined)
        else:
            raise ValueError(f"诊断模板无占位符: {tpl}")
        return q, a

    # --- 选择题（单/多选，A/B/C/D） ---
    def _build_choice_single(self, true_label: str, negs: List[str]) -> Tuple[str, str]:
        """
        单病→单选（1 正 + 3 负）。返回 (question, answer)
        答案格式： 'X {疾病名}' ，如 'A {糖尿病视网膜病变}'
        """
        if len(negs) < 3:
            raise ValueError(f"单选构造失败：负样本不足 need=3, got={len(negs)}")
        options = [true_label] + negs[:3]
        self.rng.shuffle(options)
        letters = ["A", "B", "C", "D"]
        options_str = "\n".join([f"{letters[i]}. {options[i]}" for i in range(4)])
        q = self._rc(self.choice_q_tpls).format(options=options_str)
        correct_idx = options.index(true_label)
        a = f"{letters[correct_idx]} {{{true_label}}}"
        return q, a

    def _build_choice_multi(self, true_list: List[str], negs: List[str]) -> Tuple[str, str]:
        """
        多病(<4)→多选。将 true_list 与 negs 组合成 4 选项。
        答案格式（优化后）：先给出所有正确选项字母（无逗号/无分隔），
        然后给一句简短说明，并用空格列出“字母+疾病名”的对应关系。
          例： "AC 正确。对应：A {黄斑裂孔} C {近视盘旁萎缩}"
        """
        if not (1 <= len(true_list) <= 3):
            raise ValueError(f"多选构造失败：len(true_list) 必须在 1..3，got={len(true_list)}")
        total_needed = 4
        if len(true_list) >= total_needed:
            raise ValueError("多选构造失败：true_list 数量≥4，规则要求不构造选择题。")
        need_neg = total_needed - len(true_list)
        if len(negs) < need_neg:
            raise ValueError(f"多选构造失败：负样本不足 need={need_neg}, got={len(negs)}")
    
        # 组装并打乱选项
        options = list(true_list) + negs[:need_neg]
        self.rng.shuffle(options)
    
        letters = ["A", "B", "C", "D"]
        options_str = "\n".join([f"{letters[i]}. {options[i]}" for i in range(4)])
        q = self._rc(self.choice_q_tpls).format(options=options_str)
    
        # 收集正确项的字母（按 A/B/C/D 原始顺序）
        true_set = set(true_list)
        ans_pairs = []
        ans_letters = []
        for i, opt in enumerate(options):
            if opt in true_set:
                ans_letters.append(letters[i])
                ans_pairs.append(f"{letters[i]} {{{opt}}}")
    
        # 字母连写，不加逗号；说明句子也不使用逗号分隔（用空格分隔对）
        letters_concat = "".join(ans_letters)  # 如 "AC"
        explain = "，".join([f"{letters[i]} {opt}" for i, opt in enumerate(options) if opt in true_set])
        a = f"{letters_concat} 正确。对应选项：{explain}。"
        return q, a
    
    # ========================= 采样逻辑扩展 ========================= #
    def _sample_negatives_with_relation(
        self,
        target_list: List[str],
        need: int,
        weak_ratio: float = 0.0
    ) -> List[str]:
        """
        负样本采样（固定配比版）：
          - 黑名单：绝对排除；
          - 白名单：强负样本池；
          - 其余疾病：弱负样本池；
          - 数量固定：强负=⌊need*(1-weak_ratio)⌋，弱负=need-强负；
          - 若某池不足，再从另一池补足；仍不允许采到黑名单/自身/重复。
        """
        # Fail-fast（不再吞错）：任何异常都抛出给上层
        # 1) 汇总黑/白名单
        black_set, white_set = set(), set()
        for t in target_list:
            rel = self.diease_relation_map.get(t, {})
            black_set.update(rel.get("black_list", []))
            white_set.update(rel.get("white_list", []))
        white_set -= black_set  # 黑名单优先

        # 2) 构建采样池（排除自身与黑名单）
        exclude = set(target_list) | black_set
        pool = [d for d in self.all_cn_diseases if d not in exclude]
        if not pool:
            raise ValueError(f"负样本池为空。target={target_list}, black={black_set}")

        strong_pool = [d for d in pool if d in white_set]     # 强负：白名单内
        weak_pool   = [d for d in pool if d not in white_set] # 弱负：非白名单

        # 3) 固定数量：强负=⌊need*(1-weak_ratio)⌋，弱负=need-强负
        n_strong = int(need * (1.0 - weak_ratio))
        n_weak   = need - n_strong

        # 4) 先各自抽取
        results: List[str] = []
        if n_strong > 0 and strong_pool:
            k = min(n_strong, len(strong_pool))
            results.extend(self.rng.sample(strong_pool, k=k))
        if n_weak > 0 and weak_pool:
            k = min(n_weak, len(weak_pool))
            weak_candidates = [d for d in weak_pool if d not in results]
            k = min(k, len(weak_candidates))
            if k > 0:
                results.extend(self.rng.sample(weak_candidates, k=k))

        # 5) 若还不足（某池不够），再从“另一池”补足
        remaining = need - len(results)
        if remaining > 0:
            leftovers_strong = [d for d in strong_pool if d not in results]
            take = min(remaining, len(leftovers_strong))
            if take > 0:
                results.extend(self.rng.sample(leftovers_strong, k=take))
                remaining -= take

        if remaining > 0:
            leftovers_weak = [d for d in weak_pool if d not in results]
            take = min(remaining, len(leftovers_weak))
            if take > 0:
                results.extend(self.rng.sample(leftovers_weak, k=take))
                remaining -= take

        # 6) 兜底：仍不足就从 pool 剩余里补（仍需排重）
        if remaining > 0:
            leftovers_all = [d for d in pool if d not in results]
            take = min(remaining, len(leftovers_all))
            if take > 0:
                results.extend(self.rng.sample(leftovers_all, k=take))

        return results[:need]

    # ========================= 原有工具函数 ========================= #

    def _validate_true_in_all(self, t_list: List[str]) -> None:
        not_in = [t for t in t_list if t not in self.all_cn_diseases]
        if not_in:
            raise ValueError(f"true_diag 中存在未登记的疾病: {not_in}；请先在 all_cn_diseases 中补齐。")

    def _sample_negatives_exact(self, need: int, exclude: set) -> List[str]:
        pool = [d for d in self.all_cn_diseases if d not in exclude]
        if len(pool) < need:
            raise ValueError(
                f"负样本采样不足: 需求 {need}，可用 {len(pool)}。请检查 all_cn_diseases 或 true_diag 是否过大。"
            )
        return self.rng.sample(pool, k=need)

    def _sample_two_from_list(self, arr: List[str]) -> Tuple[str, str]:
        if len(arr) < 2:
            raise ValueError("采样 2 个元素失败: 候选不足。")
        a, b = self.rng.sample(arr, k=2)
        return a, b

    def _pick_one_diagnosis(self, prompts: List[Tuple[str, str, str]]) -> Tuple[str, str, str]:
        cand = [p for p in prompts if p[2] == "diagnosis"]
        if not cand:
            raise ValueError("诊断构造中，对于质量和标签同时存在的数据无法构造问答对，请检查。")
        return self._rc(cand)

    @staticmethod
    def _dedup_preserve(items: List[str]) -> List[str]:
        seen, out = set(), []
        for x in items:
            if x not in seen:
                out.append(x)
                seen.add(x)
        return out

    def _rc(self, seq):
        """安全随机选择"""
        if isinstance(seq, list) and len(seq) == 0:
            raise ValueError("随机选择失败: 模板/候选列表为空。")
        return self.rng.choice(seq)
